import ast
import os

import matplotlib.pyplot as plt


def write_array_to_file(array, file_path,
                        *,
                        truncate=False):
    r'''
    # write array to file
    - input: array, txt file path
    '''
    folder_path = os.path.dirname(file_path)
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)

    with open(file_path, 'w+') as file:
        if truncate:
            file.truncate()
        for item in array:
            s = '[' + ', '.join(map(str, item)) + ']'
            file.write(s + '\n')


def read_array_from_file(file_path):
    r'''
    # read array from file
    - input: txt file path
    - output: array (int or float)
    '''
    array = []
    with open(file_path, 'r') as file:
        for line in file:
            line = line.strip()  # 移除行尾的换行符和空格
            if line:
                array.append(ast.literal_eval(line))
    return array


def plot_loss_accuracy_curves(loss_curves, accuracy_curves, epochs):
    """
    绘制训练过程中的 loss 和 accuracy 曲线。

    :param loss_curves: 一个包含多个 loss 曲线的列表。
    :param accuracy_curves: 一个包含多个 accuracy 曲线的列表。
    :param epochs: 一个表示 epochs 的整数。
    """
    fig, ax1 = plt.subplots()

    # 绘制 loss 曲线
    color = 'tab:red'
    ax1.set_xlabel('Epoch')
    ax1.set_ylabel('Loss', color=color)
    for i, curve in enumerate(loss_curves):
        curve = [p for p in curve]
        ax1.plot(range(0, epochs), curve, color=color, label='Validation Loss')
    ax1.tick_params(axis='y', labelcolor=color)

    # 实例化第二个坐标轴，共享同一个 x 轴
    ax2 = ax1.twinx()
    color = 'tab:blue'
    ax2.set_ylabel('Accuracy', color=color)
    # ax2.set_ylim([0, 1])  # 设置 accuracy 的 y 轴范围为 0 到 1

    colors = ['tab:blue', 'tab:green', 'tab:orange']
    labels = ['Train Accuracy', 'Validation Accuracy', 'Test Accuracy']

    for curve, color, label in zip(accuracy_curves, colors, labels):
        curve = [p for p in curve]
        ax2.plot(range(0, epochs), curve, color=color, label=label, linestyle='dashed')

    ax2.tick_params(axis='y', labelcolor=color)

    lines, labels = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax2.legend(lines + lines2, labels + labels2, loc='center right')

    fig.tight_layout()
    plt.title("Training Loss and Accuracy Curves")
    plt.show()
